"""
步态生成器
用于生成六足机器人的各种步态模式
"""

import math
import numpy as np
from typing import List, Dict, Tuple
from dataclasses import dataclass


@dataclass
class GaitParams:
    """步态参数"""

    stride_length: float = 50.0  # 步长
    stride_height: float = 30.0  # 抬腿高度
    cycle_time: float = 1.0  # 步态周期
    duty_factor: float = 0.5  # 支撑相占比
    body_height: float = 100.0  # 身体高度
    body_roll: float = 0.0  # 身体横滚角
    body_pitch: float = 0.0  # 身体俯仰角
    body_yaw: float = 0.0  # 身体偏航角


class GaitGenerator:
    """步态生成器类"""

    def __init__(self, leg_config: Dict[int, Tuple[float, float, float]]):
        """
        初始化步态生成器

        Args:
            leg_config: 腿部配置，格式为 {leg_id: (x, y, z)}
        """
        self.leg_config = leg_config
        self.params = GaitParams()

    def tripod_gait(
        self, t: float, direction: Tuple[float, float] = (1.0, 0.0)
    ) -> Dict[int, Tuple[float, float, float]]:
        """
        三脚架步态

        Args:
            t: 时间参数 [0, 1]
            direction: 运动方向 (x, y)

        Returns:
            各腿的目标位置 {leg_id: (x, y, z)}
        """
        positions = {}

        # 定义三脚架分组
        tripod1 = [1, 3, 5]  # 第一组腿
        tripod2 = [2, 4, 6]  # 第二组腿

        # 计算摆动相和支撑相
        swing_phase = t % 1.0
        is_swing1 = swing_phase < self.params.duty_factor
        is_swing2 = not is_swing1

        for leg_id in self.leg_config:
            base_pos = self.leg_config[leg_id]

            if (leg_id in tripod1 and is_swing1) or (leg_id in tripod2 and is_swing2):
                # 摆动相
                phase = (
                    swing_phase / self.params.duty_factor
                    if leg_id in tripod1
                    else (swing_phase - self.params.duty_factor)
                    / (1 - self.params.duty_factor)
                )
                x = base_pos[0] + self.params.stride_length * direction[0] * phase
                y = base_pos[1] + self.params.stride_length * direction[1] * phase
                z = base_pos[2] + self.params.stride_height * math.sin(phase * math.pi)
            else:
                # 支撑相
                phase = (
                    (swing_phase - self.params.duty_factor)
                    / (1 - self.params.duty_factor)
                    if leg_id in tripod1
                    else swing_phase / self.params.duty_factor
                )
                x = base_pos[0] - self.params.stride_length * direction[0] * phase
                y = base_pos[1] - self.params.stride_length * direction[1] * phase
                z = base_pos[2]

            positions[leg_id] = (x, y, z)

        return positions

    def wave_gait(
        self, t: float, direction: Tuple[float, float] = (1.0, 0.0)
    ) -> Dict[int, Tuple[float, float, float]]:
        """
        波浪步态

        Args:
            t: 时间参数 [0, 1]
            direction: 运动方向 (x, y)

        Returns:
            各腿的目标位置 {leg_id: (x, y, z)}
        """
        positions = {}

        # 定义腿的相位差
        phase_offsets = {1: 0.0, 2: 0.2, 3: 0.4, 4: 0.6, 5: 0.8, 6: 1.0}

        for leg_id in self.leg_config:
            base_pos = self.leg_config[leg_id]
            phase = (t + phase_offsets[leg_id]) % 1.0

            if phase < self.params.duty_factor:
                # 支撑相
                x = (
                    base_pos[0]
                    - self.params.stride_length
                    * direction[0]
                    * phase
                    / self.params.duty_factor
                )
                y = (
                    base_pos[1]
                    - self.params.stride_length
                    * direction[1]
                    * phase
                    / self.params.duty_factor
                )
                z = base_pos[2]
            else:
                # 摆动相
                swing_phase = (phase - self.params.duty_factor) / (
                    1 - self.params.duty_factor
                )
                x = base_pos[0] + self.params.stride_length * direction[0] * swing_phase
                y = base_pos[1] + self.params.stride_length * direction[1] * swing_phase
                z = base_pos[2] + self.params.stride_height * math.sin(
                    swing_phase * math.pi
                )

            positions[leg_id] = (x, y, z)

        return positions

    def ripple_gait(
        self, t: float, direction: Tuple[float, float] = (1.0, 0.0)
    ) -> Dict[int, Tuple[float, float, float]]:
        """
        涟漪步态

        Args:
            t: 时间参数 [0, 1]
            direction: 运动方向 (x, y)

        Returns:
            各腿的目标位置 {leg_id: (x, y, z)}
        """
        positions = {}

        # 定义腿的相位差
        phase_offsets = {1: 0.0, 2: 0.4, 3: 0.2, 4: 0.6, 5: 0.4, 6: 0.8}

        for leg_id in self.leg_config:
            base_pos = self.leg_config[leg_id]
            phase = (t + phase_offsets[leg_id]) % 1.0

            if phase < self.params.duty_factor:
                # 支撑相
                x = (
                    base_pos[0]
                    - self.params.stride_length
                    * direction[0]
                    * phase
                    / self.params.duty_factor
                )
                y = (
                    base_pos[1]
                    - self.params.stride_length
                    * direction[1]
                    * phase
                    / self.params.duty_factor
                )
                z = base_pos[2]
            else:
                # 摆动相
                swing_phase = (phase - self.params.duty_factor) / (
                    1 - self.params.duty_factor
                )
                x = base_pos[0] + self.params.stride_length * direction[0] * swing_phase
                y = base_pos[1] + self.params.stride_length * direction[1] * swing_phase
                z = base_pos[2] + self.params.stride_height * math.sin(
                    swing_phase * math.pi
                )

            positions[leg_id] = (x, y, z)

        return positions

    def set_params(self, params: GaitParams) -> None:
        """设置步态参数"""
        self.params = params
